// https://alist.home.easyfind.top:36666
package main

import (
	"bufio"
	"flag"
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"

	"github.com/rs/zerolog"

	"time"
)

var logger = zerolog.New(os.Stdout).With().Timestamp().Logger()

func main() {
	help := flag.Bool("help", false, "print usage")
	bind := flag.String("bind", "0.0.0.0:9528", "The address to bind to")
	backend := flag.String("backend", "4.easyfind.top:35001", "The backend server address")
	flag.Parse()

	logger.Level(zerolog.DebugLevel)

	if *help {
		flag.Usage()
		return
	}

	if *backend == "" {
		flag.Usage()
		return
	}

	if *bind == "" {
		//use default bind
		logger.Info().Str("bind", *bind).Msg("use default bind")
	}

	// create log dir
	err := os.MkdirAll("logs", 0755)
	if err != nil {
		logger.Error().Err(err).Send()
		os.Exit(1)
	}

	success, err := RunProxy(*bind, *backend)
	if !success {
		logger.Error().Err(err).Send()
		os.Exit(1)
	}
}

func RunProxy(bind, backend string) (bool, error) {
	listener, err := net.Listen("tcp", bind)
	if err != nil {
		return false, err
	}
	defer listener.Close()
	logger.Info().Str("bind", bind).Str("backend", backend).Msg("tcp-proxy started.")
	for {
		conn, err := listener.Accept()
		if err != nil {
			logger.Error().Err(err).Send()
		} else {
			go ConnectionHandler(conn, backend)
		}
	}
}

func ConnectionHandler(conn net.Conn, backend string) {
	logger.Info().Str("conn", conn.RemoteAddr().String()).Msg("client connected.")
	target, err := net.Dial("tcp", backend)
	defer conn.Close()
	if err != nil {
		logger.Error().Err(err).Send()
	} else {
		defer target.Close()
		logger.Info().Str("conn", conn.RemoteAddr().String()).Str("backend", target.LocalAddr().String()).Msg("backend connected.")
		closed := make(chan bool, 1)

		go Proxy(conn, target, closed)
		go Proxy2(target, conn, closed)
		<-closed
		logger.Info().Str("conn", conn.RemoteAddr().String()).Msg("Connection closed.")
	}
}

func Proxy(from net.Conn, to net.Conn, closed chan bool) {

	buffer := make([]byte, 4096)
	for {
		n1, err := from.Read(buffer)
		if err != nil {
			closed <- true
			return
		}
		fmt.Println("===>read:", n1)
		fmt.Println("===>data:", string(buffer[:n1]))
		n2, err := to.Write(buffer[:n1])
		logger.Debug().Str("from", from.RemoteAddr().String()).Int("recv", n1).Str("to", to.RemoteAddr().String()).Int("send", n2).Send()
		if err != nil {
			closed <- true
			return
		}
	}

}

func Proxy2(from net.Conn, to net.Conn, closed chan bool) {
	time.Sleep(1 * time.Second)
	buffer := make([]byte, 4096)

	// Open a log file for writing
	logFile, err := os.OpenFile("logs/proxy.log", os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		fmt.Println("Failed to open log file:", err)
		closed <- true
		return
	}
	defer logFile.Close()

	for {
		n1, err := from.Read(buffer)
		if err != nil {
			closed <- true
			return
		}
		fmt.Println("<===read:", n1)
		fmt.Println("<===data:", string(buffer[:n1]))
		n2, err := to.Write(buffer[:n1])
		logger.Debug().Str("from", from.RemoteAddr().String()).Int("recv", n1).Str("to", to.RemoteAddr().String()).Int("send", n2).Send()
		if err != nil {
			closed <- true
			return
		}
		// Write received data to the log file
		if _, err := logFile.Write(buffer[:n1]); err != nil {
			fmt.Println("Failed to write to log file:", err)
			return
		}

	}
}

// 解析HTTP报文
func parseHttp(conn net.Conn) string {

	var s = ""
	reader := bufio.NewReader(conn)
	firstLine, _, _ := reader.ReadLine()
	s += string(firstLine) + "\n"

	tempList := strings.Split(string(firstLine), " ")
	method := tempList[0]
	path := tempList[1]
	proto := tempList[2]
	fmt.Println(1111, method, path, proto)
	var headers = make(map[string]string)
	var contenLength int
	for {
		line, _, _ := reader.ReadLine()

		s += string(line) + "\n"
		if len(line) == 0 {
			break
		}
		tempList = strings.Split(string(line), ":")
		key := tempList[0]
		value := strings.Trim(tempList[1], " ")
		headers[key] = value
		if key == "Content-Length" {
			contenLength, _ = strconv.Atoi(value)
		}
	}

	fmt.Println(2222, headers)
	if contenLength == 0 {
		return s
	}
	pack := make([]byte, contenLength)
	t, _ := reader.Read(pack)
	s += string(pack[:t])

	return s
}
